import sys
sys.path.insert(-1, "..")
from llm.doubao_demo import DoubaoClient
from utils.yj_api import YJAPI
from utils.cv2_utils import get_image_from_url, pillow_to_cv2
import json
from components.xiaokuai_detect import get_xiaokuai_result
from components.yolo_api import YoloAPI

class StructureExtractor:
    def __init__(self, doubao_api_key=None, doubao_base_url=None):
        self.doubao_client = DoubaoClient(doubao_api_key, doubao_base_url)
        self.yj_api = YJAPI()
        self.yolo_api = YoloAPI()
        # 是否检测插图？默认开启
        self.to_detect_figure = True        

    def extract_structure(self, subject_id):
        """根据subject_id提取图片结构
        
        Args:
            subject_id (str): 题目ID
            
        Returns:
            dict: 包含图片结构信息的字典
        """
        try:
            # 获取图片URL
            image_urls = self.yj_api.get_image_urls(subject_id)

            print("image url : ", image_urls)
            
            if not image_urls:
                print("error : 未找到相关图片")
                return None
                
            # 构建提示词
            text = '''这是一份考试的答题卡图片，请对图片进行OCR，并提取里面的试卷标题和题目结构。试卷标题包括大题号、小题号、分值。如果是选择题（客观题），除了给出题号外，还需要给出该小题的选项数量。请以json的形式输出试卷的题目结构，无需详细的题干内容。输出的json格式如下：{"title":"八年级教学质量监测数学答题卡","sections":[{"sectionNumber":"一","sectionTitle":"选择题","score":24,"subQuestions":[{"questionNumber":1,"optionCount":4}]},{"sectionNumber":"二","sectionTitle":"填空题","score":15,"subQuestions":[{"questionNumber":9}]},{"sectionNumber":"三","sectionTitle":"解答题","score":81,"subQuestions":[{"questionNumber":20,"score":5},{"questionNumber":21,"score":6,"subSubQuestions":[{"questionNumber":"(1)"},{"questionNumber":"(2)"}]}]}]}'''
            
            # 构建消息
            messages = self.doubao_client.build_message(text, image_urls)
            
            # 调用大模型
            result = self.doubao_client.call_model("ep-20250124123709-scsq6", messages)

            print("result : ", result)
            
            # 生成大纲
            outline = self.doubao_client.print_outline(result)
            
            # 返回结构信息
            return {
                "subject_id": subject_id,
                "outline":outline, # 返回大纲字符串
                "result": result,  
                "image_urls": image_urls  # 返回图片URL列表
            }
        except Exception as e:
            print(f"Error extracting structure: {str(e)}")
            return None
        

    def extract_structure_from_image_urls(self, image_urls, image_type="dtk"):
        """根据subject_id提取图片结构
        
        Args:
            image_urls (list): 图片链接
            
        Returns:
            dict: 包含图片结构信息的字典
        """
        try:
            # print("image url : ", image_urls)
            
            if not image_urls:
                print("error : 未找到相关图片")
                return None
                
            # 构建提示词
            text = '''这是一份考试的答题卡图片，请对图片进行OCR，并提取里面的试卷标题和题目结构。试卷标题包括大题号、小题号、分值。如果是选择题（客观题），除了给出题号外，还需要给出该小题的选项数量。请以json的形式输出试卷的题目结构，无需详细的题干内容。输出的json格式如下：{"title":"八年级教学质量监测数学答题卡","sections":[{"sectionNumber":"一","sectionTitle":"选择题","score":24,"subQuestions":[{"questionNumber":1,"optionCount":4}]},{"sectionNumber":"二","sectionTitle":"填空题","score":15,"subQuestions":[{"questionNumber":9}]},{"sectionNumber":"三","sectionTitle":"解答题","score":81,"subQuestions":[{"questionNumber":20,"score":5},{"questionNumber":21,"score":6,"subSubQuestions":[{"questionNumber":"(1)"},{"questionNumber":"(2)"}]}]}]}'''
            
            # 构建消息
            messages = self.doubao_client.build_message(text, image_urls)
            # print("messages : ", messages)
            
            # 调用大模型
            result = self.doubao_client.call_model("ep-20250124123709-scsq6", messages)

            # print("result : ", result)

            if image_type == "dtk":
                result = self._fix_llm_result(result, image_urls)

            # print("fixed result : ", result)

            # 检测图片中的插图
            if self.to_detect_figure:
                print("begin to call yolo api to detect figures...")
                all_figures = []
                for img_url in image_urls:
                    pil_img = get_image_from_url(img_url)
                    figures = []
                    if pil_img is not None:
                        cv_image = pillow_to_cv2(pil_img)
                        if cv_image is not None:
                            figures = self.yolo_api.analyse_layout_of_cv_image(cv_image, select_types=['figure'])
                    all_figures.append(figures)
                # 添加到结果中
                result['figures'] = all_figures
            
            # 生成大纲
            outline = self.doubao_client.print_outline(result)
            
            # 返回结构信息
            return {
                "subject_id": None,
                "outline":outline, # 返回大纲字符串
                "result": result,  
                "image_urls": image_urls  # 返回图片URL列表
            }
        except Exception as e:
            print(f"Error extracting structure: {str(e)}")
            return None

    def extract_structure_from_local_file(self, subject_id, result_file):
        """根据subject_id提取图片结构
        
        Args:
            subject_id (str): 题目ID
            
        Returns:
            dict: 包含图片结构信息的字典
        """
        try:
            # 获取图片URL
            image_urls = self.yj_api.get_image_urls(subject_id)

            print("image url : ", image_urls)
            
            if not image_urls:
                print("error : 未找到相关图片")
                return None
                
            # # 构建提示词
            # text = '''这是一份考试的答题卡图片，请对图片进行OCR，并提取里面的试卷标题和题目结构。试卷标题包括大题号、小题号、分值。如果是选择题（客观题），除了给出题号外，还需要给出该小题的选项数量。请以json的形式输出试卷的题目结构，无需详细的题干内容。输出的json格式如下：{"title":"八年级教学质量监测数学答题卡","sections":[{"sectionNumber":"一","sectionTitle":"选择题","score":24,"subQuestions":[{"questionNumber":1,"optionCount":4}]},{"sectionNumber":"二","sectionTitle":"填空题","score":15,"subQuestions":[{"questionNumber":9}]},{"sectionNumber":"三","sectionTitle":"解答题","score":81,"subQuestions":[{"questionNumber":20,"score":5},{"questionNumber":21,"score":6,"subSubQuestions":[{"questionNumber":"(1)"},{"questionNumber":"(2)"}]}]}]}'''
            
            # # 构建消息
            # messages = self.doubao_client.build_message(text, image_urls)
            
            # # 调用大模型
            # result = self.doubao_client.call_model("ep-20250124123709-scsq6", messages)

            with open(result_file, 'r', encoding='utf-8') as fin:
                result = json.load(fin)


            print("result : ", result)
            result["image_urls"] = image_urls
            return result
        except Exception as e:
            print(f"Error extracting structure: {str(e)}")
            return None

    def save_structure(self, structure, output_path):
        """保存提取的结构到文件
        
        Args:
            structure (dict): 结构信息
            output_path (str): 输出文件路径
        """
        with open(output_path, 'w', encoding='utf-8') as f:
            json.dump(structure, f, ensure_ascii=False, indent=2)

    def print_structure(self, structure):
        """打印提取的结构信息
        
        Args:
            structure (dict): 结构信息
        """
        print(json.dumps(structure, ensure_ascii=False, indent=2))

    def print_outline(self, structure):
        """生成试卷结构大纲
        
        Args:
            structure (dict): 结构信息
            
        Returns:
            str: 试卷结构大纲字符串
        """
        return self.doubao_client.print_outline(structure.get('structure', {}))
    

    
    def _get_objective_question(self, result):
        '''
        从大模型提取的结果中找出客观题
        '''

        objective_questions = []

        if result is None:
            return objective_questions       
        
        if "sections" in result:

            for section in result.get('sections', []):                
                for question in section.get('subQuestions', []):
                    q_num = question.get('questionNumber', '')
                    score = question.get('score', '')
                    options = question.get('optionCount', '')
                    
                    if options:
                        objective_questions.append((q_num, options))
        
        return objective_questions
    
    def _fix_llm_result(self, llm_result, image_urls):
        '''
        使用原本的检测模型（小块检测服务）的结果来修复大模型可能的有误的地方
        从测试来看，原本的检测模型对客观题效果较好，大模型对主观题较好
        '''

        objective_questions_from_llm = self._get_objective_question(llm_result)

        if len(objective_questions_from_llm) > 0:
            # 存在客观题，则去下载图片及调小块检测的服务
            # 使用小块检测服务的客观题选项块数来验证或修复大模型客观题选项块数
            # 一般来说客观题在第一页，只取第一张图去检测

            url_of_first_image = image_urls[0]
            pil_img = get_image_from_url(url_of_first_image)
            if pil_img is not None:

                cv_image = pillow_to_cv2(pil_img)

                if cv_image is not None:

                    xiaokuai_result, _ = get_xiaokuai_result(cv_image, "", "", 0)
                    objective_questions_from_detect_model = xiaokuai_result.get_obj_list()

                    # 遍历大模型提取的客观题，与小块检测的结果做对比，找出选项数有误的客观题
                    wrong_option_questions = []
                    for question in objective_questions_from_llm:

                        q_num, options = question

                        for item in objective_questions_from_detect_model:
                            q_num_1, options_1 = item
                            if q_num == q_num_1:
                                if options != options_1:
                                    # 改为正确的选项数
                                    wrong_option_questions.append((q_num, options_1))
                                    break
                    
                    # 修复选项数
                    if len(wrong_option_questions) > 0:
                        llm_result = self._fix_wrong_objective_questions(wrong_option_questions, llm_result)
        return llm_result

    
    def _fix_wrong_objective_questions(self, wrong_option_questions, result):
        '''
        根据找出来的选项数错误的客观题及原始的大模型结果，修复大模型结果
        '''

        print("\nwrong_option_questions : ", wrong_option_questions)

        for question in wrong_option_questions:
            q_num, options = question

            for s_index, section in enumerate(result.get('sections', [])):
                found_question = False
                for q_index, question in enumerate(section.get('subQuestions', [])):
                    q_num_1 = question.get('questionNumber', '')
                    options_1 = question.get('optionCount', '')
                    
                    if options_1 and q_num_1 == q_num:
                        result['sections'][s_index]['subQuestions'][q_index]['optionCount'] = options
                        found_question = True
                        break
                if found_question:
                    break
        
        return result

# 测试代码
if __name__ == "__main__":
    import asyncio
    
    async def main():
        extractor = StructureExtractor()
        
        # 测试subject_id
        subject_id = "9361325"
        # # 提取
        # structure = await extractor.extract_structure(subject_id)

        image_urls = [
        "https://yj-oss.yunxiao.com/v1/baidu-raw/template/9361325/F2A99459-CD6A-11EF-9907-0C96E69DE468.png?authorization=bce-auth-v1%2Fa908715249bb41c998c7d924b2476b37%2F2025-02-06T03%3A13%3A09Z%2F3600%2Fhost%2Fa3ef5b5fd1d79b1fe470c3e9e88b1f8993335ee0054caca57dfd8d6278b166a6",
        "https://yj-oss.yunxiao.com/v1/baidu-raw/template/9361325/F2A9945A-CD6A-11EF-9907-0C96E69DE468.png?authorization=bce-auth-v1%2Fa908715249bb41c998c7d924b2476b37%2F2025-02-06T03%3A13%3A09Z%2F3600%2Fhost%2F93d23c6198dfd93e856e2211d0c0c09e152e2c220bf6e39317d6484efe5d7b4c"
    ]
        
        # 提取
        structure = await extractor.extract_structure_from_image_urls(image_urls)
        if structure is not None:
            # 保存结构到文件
            extractor.save_structure(structure, f"./{subject_id}_structure.json")
            
            # 打印试卷大纲
            outline = structure['result']
            print(f"\n科目ID : {subject_id} 试卷大纲：")
            print(outline)
    
    asyncio.run(main())
